package t20240902_udf;

import com.maxmind.db.Reader;
import org.apache.flink.table.functions.FunctionContext;
import org.apache.flink.table.functions.ScalarFunction;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.InetAddress;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class ParseIpPro extends ScalarFunction
{

    public static String IPV4_IPDB = "D:\\code\\flink-all\\flink_demo\\ipdb_enc_ipv4.mmdb";
    protected static final String SECRET_KEY_FILE = "D:\\code\\flink-all\\flink_demo\\secret.key";

    private Reader readerV4 = null;

    private long lastMidifyTime = 0L;
    private Timer timer;


    private long itd;

    private ReentrantReadWriteLock lock = new ReentrantReadWriteLock();

    public static String[] parseSecretKey(FileSystem fs) throws IOException {
        ByteArrayOutputStream outSteam = null;
        String[] tempStr = null;
        InputStream is = null;

        try {
            is = fs.open(new Path(SECRET_KEY_FILE));
            outSteam = new ByteArrayOutputStream();
            byte[] buffer = new byte[1024];
            int len;
            while((len = is.read(buffer)) != -1) {
                outSteam.write(buffer, 0, len);
            }

            tempStr = outSteam.toString().split("@");
        } catch (Exception var9) {
        } finally {

            if (outSteam != null) {
                outSteam.close();
            }
            if (is != null) {
                is.close();
            }
        }

        return tempStr;
    }

    @Override
    public void open(FunctionContext context) throws Exception {
        itd = Thread.currentThread().getId();
        System.out.println(itd + " open run ");
        FileSystem fs = FileSystem.get(new Configuration());
        flushReader(fs);
        timer = new Timer();
        timer.scheduleAtFixedRate(new TimerTask() {
            @Override
            public void run() {
                reloadMMDbFile(fs);
            }
        }, 60, 60 * 1000);
    }

    private void reloadMMDbFile(FileSystem fs) {
        try {
            FileStatus fileStatus = fs.getFileStatus(new Path(IPV4_IPDB));
            long modificationTime = fileStatus.getModificationTime();
            if (modificationTime != lastMidifyTime) {
                flushReader(fs);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private void flushReader(FileSystem fs) {
        FSDataInputStream fsisV4 = null;
        InputStream cipherInputStreamV4 = null;
        try {
            String[] secret = IpHelper.parseSecretKey(fs);
            fsisV4 = fs.open(new Path(IPV4_IPDB));
            cipherInputStreamV4 = AES256Util.doDecrypt(secret[0], secret[1], fsisV4);

            lock.writeLock().lock();
            try {
                if (readerV4 != null) {
                    readerV4.close();
                }
                readerV4 = new Reader(cipherInputStreamV4);
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                lock.writeLock().unlock();
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            IpHelper.closeStream(fsisV4);
            IpHelper.closeStream(cipherInputStreamV4);
        }
    }

    public String eval(String ip){
        if (itd != Thread.currentThread().getId()) {
            throw new RuntimeException("Thread Wrong！！！！");
        }
        lock.readLock().lock();
        try {
            byte[] addr = InetAddress.getByName(ip).getAddress();
            InetAddress address = InetAddress.getByAddress(addr);
            return IpHelper.generateIpInfo(address, readerV4);
        }catch (Exception ex){
            ex.printStackTrace();
        } finally {
            lock.readLock().unlock();
        }
        return IpHelper.NULL_IP_CODE;
    }

    @Override
    public void close() throws Exception {
        System.out.println("close run ");
        if (timer != null) {
            timer.cancel();
        }

    }
}
